package dd.lo.classifier;

import java.io.BufferedInputStream;
import java.io.Closeable;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.zip.GZIPInputStream;

/**
 * 关于MNIST数据集的单文件数据排列方式定义讲解
 * https://blog.51cto.com/u_14850/7707079
 */
public class MnistReader implements Closeable {

    /**
     * 从字节转整型参考：
     * https://www.cnblogs.com/kendoziyu/p/15575059.html
     * 需要注意字节是带符号的，强转整型会有负数，需要先利用&0xFF拓展为补码
     */
    private static int fourByteArr2Integer(byte[] nBuf) {
        if (4 != nBuf.length) throw new IllegalStateException("buf不是4个字节长度！");
        return ((nBuf[0] & 0xFF) << 24) | ((nBuf[1] & 0xFF) << 16) | ((nBuf[2] & 0xFF) << 8) | (nBuf[3] & 0xFF);
    }

    public static int TYPE_IMAGE = 0;
    public static int TYPE_LABEL = 1;

    private final GZIPInputStream inputStream;

    private final int n;

    private final int height;

    private final int width;

    private final List<byte[]> dataList;

    public MnistReader(String filePath, int type) throws IOException {
        if (type != TYPE_IMAGE && type != TYPE_LABEL) throw new IllegalArgumentException("必须指定读取的是图片数据集还是标签数据集！");
        //GZIPInputStream的size参数是缓冲区大小，与read函数有关，如果缓冲区较小，则read会较多出现读取字节数不足目标数字大小的情况，需要追加读取补充完整
        inputStream = new GZIPInputStream(new BufferedInputStream(Files.newInputStream(Paths.get(filePath))), 8192);
        //Magic number
        byte[] magicNumber = new byte[4];
        int readLen = inputStream.read(magicNumber);
        if (readLen != 4) throw new IllegalStateException("Magic number 读取失败！");
        int dimension = magicNumber[3];
        System.out.printf("Dimension: %d\n", dimension);
        if (dimension == 1 && type != TYPE_LABEL) throw new IllegalStateException("此文件应该是标签数据集，type错误");
        //N
        byte[] nBuf = new byte[4];
        readLen = inputStream.read(nBuf);
        if (readLen != 4) throw new IllegalStateException("样本数读取失败！");
        n = fourByteArr2Integer(nBuf);
        if (type == TYPE_IMAGE) {
            //Height
            byte[] hBuf = new byte[4];
            readLen = inputStream.read(hBuf);
            if (readLen != 4) throw new IllegalStateException("单个样本行数读取失败！");
            height = fourByteArr2Integer(hBuf);
            //Width
            byte[] wBuf = new byte[4];
            readLen = inputStream.read(wBuf);
            if (readLen != 4) throw new IllegalStateException("单个样本列数读取失败！");
            width = fourByteArr2Integer(wBuf);
        } else {
            height = 1;
            width = 1;
        }
        int i = 0;
        int dataSize = height * width;
        dataList = new ArrayList<>(n);
        while (i < n) {
            byte[] data = new byte[dataSize];
            readLen = inputStream.read(data);
            if (readLen < 0) throw new RuntimeException(String.format("第%d张图片数据读取失败，遇到文件结尾", i));
            while (readLen < dataSize) {
                //有时候读取的数据块大小不足dataSize，这是因为读取到缓冲区边缘了，这时候需要进行补充读取，并追加到data数组中
                int residue = dataSize - readLen;
                byte[] supplement = new byte[residue];
                int len = inputStream.read(supplement);
                System.arraycopy(supplement, 0, data, readLen, len);
                readLen += len;
            }
            dataList.add(data);
            ++i;
        }
    }

    public byte[] get(int idx) {
        if (idx >= n) throw new IndexOutOfBoundsException(String.format("有效范围%d~%d，但读取%d位置数据", 0, n - 1, idx));
        return dataList.get(idx);
    }

    public int size() {
        return n;
    }

    public int getHeight() {
        return height;
    }

    public int getWidth() {
        return width;
    }

    @Override
    public void close() throws IOException {
        if (null != inputStream) {
            inputStream.close();
        }
    }

    public static void main(String[] args) {
        String filePath = "/Users/kwok/Downloads/MNIST/train-images-idx3-ubyte.gz";
        try (MnistReader reader = new MnistReader(filePath, TYPE_IMAGE)) {
            byte[] bytes = reader.get(2);
            System.out.println(Arrays.toString(bytes));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
